11609. Finding pairs

 

You are given a rooted tree with n nodes and node 11 as a root. There is a unique path between any two nodes. Here, d(i, j) is defined as a number of edges in a unique path between nodes i and j.

You have to find the number of pairs (i, j) such that d(i, j) = d(i, 1) – d(j, 1).

 

Input. The first line contains the number of the nodes n (1 n 105) in the tree. Each of the next n – 1 lines denotes the edge of the tree.

 

Output. Print a single integer denoting the number of pairs (i, j) such that d(i, j) = d(i, 1) – d(j, 1).

 

Sample input

Sample output

5

1 2

2 3

2 4

4 5

13

 

SOLUTION

graphs depth first search

 

Algorithm analysis

The condition d(i, j) = d(i, 1) – d(j, 1) is valid for any pair (i, j) for which j is the ancestor of i in a depth first search from the node 1. Consider an example:

·        d(4, 1) = 2, d(4, 1) – d(1, 1) = 2 – 0 = 2;

·        d(6, 3) = 2, d(6, 1) – d(3, 1) = 3 1 = 2;

·        d(2, 2) = 0, d(2, 1) – d(2, 1) = 1 1 = 0;

·        d(6, 1) = 3, d(6, 1) – d(1, 1) = 3 0 = 3;

 

The depth h[v] of a vertex v is the number of vertices on the path from 1 to v. In the figure, the depth is indicated near each vertex. Let us fix vertex i and answer the question: how many vertices j exist such that d(i, j) = d(i, 1) – d(j, 1).

For example, for i = 6 there are 4 such vertices: j = 1, 3, 4, 6. Note that h[6] = 4. For a fixed value of i, there are exactly h[i] corresponding values of j.

To solve the problem, for each vertex v of the tree, find its depth h[v] and compute the sum of the depths over all vertices.

 

Example

Consider the graph from a sample. Near each vertex its depth is written.

The sum of the depths of all vertices is 1 + 2 + 3 + 3 + 4 = 13.

 

Algorithm realization

Declare the adjacency list of the graph g. Store the depth of the vertices in the array h.

 

vector<vector<int> > g;

vector<int> h;

 

The function dfs for each vertex v computes its depth. The parent of v is the vertex p.

 

int dfs(int v, int p)

{

  for (int to : g[v])

 

Process the edge (v, to). If to is not a parent of v, then we compute the depth of to and run a depth first search from to.

 

    if (to != p)

    {

      h[to] = h[v] + 1;

      dfs(to, v);

    }

  return h[v];

}

 

The main part of the program. Read the number of vertices n in the tree.

 

scanf("%d", &n);

g.resize(n + 1);

 

Construct a tree.

 

for (i = 2; i <= n; i++)

{

  scanf("%d %d", &a, &b);

  g[a].push_back(b);

  g[b].push_back(a);

}

 

Let the height of the vertex 1 be equal to 1.

 

h.resize(n + 1);

h[1] = 1;

 

Start a depth first search from the vertex 1.

 

dfs(1, -1);

 

Compute the answer – the sum of the depths of all vertices.

 

res = 0;

for (i = 1; i <= n; i++)

  res += h[i];

 

Print the answer.

 

printf("%lld\n", res);